# pragma once

# include <homo_ocp.h>

# include "cpp_tools.h"
# include "udpSettings.h"
# include "estimator_base.h"

# include "filter_W1M1_C.h"

# ifdef INPUT_REAL
# include "LTP_A_anchor.h"
# include "LTP_A_tag.h"
# include "WHEELTEC_N100.h"
# endif // INPUT_REAL

class estimator_UWB_IMU : public estimator_base
{
private:
    struct mtData
    {
        cpp_tools::mtDataManager<Eigen::VectorXd> IMU_filter = Eigen::VectorXd::Zero(3); // 滤波加速度数据 a_x,a_y,a_z
        cpp_tools::mtDataManager<Eigen::VectorXd> AHRS_filter = Eigen::VectorXd::Zero(3); // 滤波陀螺仪数据 psi,gamma,omega_z
        
        cpp_tools::mtDataManager<Eigen::VectorXd> uwb_front = Eigen::VectorXd::Zero(2); // x, y
        cpp_tools::mtDataManager<int> uwb_front_update = 0;
        
        cpp_tools::mtDataManager<Eigen::VectorXd> uwb_rear = Eigen::VectorXd::Zero(2); // x, y
        cpp_tools::mtDataManager<int> uwb_rear_update = 0;

        cpp_tools::mtDataManager<Eigen::VectorXd> state;

        cpp_tools::mtDataManager<std::pair<double,double>> position_cor = {0,0};
    } m_mtData;

    // =============================定义参数=============================
    struct frequency // 频率，Hz
    {
        # ifdef INPUT_REAL
        double imu_port = 1e4;
        double uwb_port = 1e4;
        # else
        double state = 1e4;
        # endif

        double estimate = 200;
    } const m_frequency;

    const double gravity = 9.82; // 重力加速度
    const double pi = cpp_tools::pi; // 圆周率

    // UWB设备
    std::string UWB_A_device = "/dev/ttyACM0";
    std::string UWB_B_device = "/dev/ttyACM1";
    int UWB_bauds = 921600;

    // UWB id
    const double front_UWB_ID = 2; // 定义前面的UWB id，后面的默认就是另一个

    // UWB到质心的距离
    const double dx_front = 0.216; // 前面的uwb到质心的距离
    const double dx_rear = -0.258; // 后面的uwb到质心的距离

    // 初始化用到的数据量
    const size_t initDataLen = 200;

    // IMU设备
    std::string IMU_device = "/dev/ttyUSB0";
    int IMU_bauds = 921600;

    // IMU滤波相关
    const int avgFilterLen = 5; // 均值滤波需要的数据长度
    const double maxDis = 5; // 加速度错误值滤波的最大限差

    // AHRS滤波相关
    const double ahrsFilterFactor = 0.995;

    // ==================================================================

    cpp_tools::threadPool m_threadPool; // 线程池

    bool isRunning = true;
            
    double t_update = -1; // 上一次状态估计更新的时间

    // ==================================================================
    /// @brief 从数据库中读取imu和uwb的数据
    /// @param imu 写入imu数据到此处 [psi,gamma,omega,a_x,a_y]
    /// @param uwb 写入uwb数据到此处 [x,y,varphi]
    /// @return -1：写入失败，有无效数据或者uwb数据没有更新（由于imu频率大于uwb，因此认为当uwb更新时imu数据总是更新后的）；
    ///          0：写入成功
    int getData(Eigen::VectorXd & imu, Eigen::VectorXd & uwb)
    {
        if (!*m_mtData.uwb_front_update.read().access())
        {
            // std::cout << "waitting for uwb A data ============================================" << std::endl;
            return -1;
        }
        if (!*m_mtData.uwb_rear_update.read().access())
        {
            // std::cout << "waitting for uwb B data >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>" << std::endl;
            return -1;
        }
        
        {
            auto ac_imu = m_mtData.IMU_filter.read();
            Eigen::VectorXd a(3);
            a << (*ac_imu.access())(0), 
            -(*ac_imu.access())(1),
            -(*ac_imu.access())(2);

            auto ac_ahrs = m_mtData.AHRS_filter.read();
            double psi = -(*ac_ahrs.access())(0);
            double gamma = (*ac_ahrs.access())(1);
            double omega = -(*ac_ahrs.access())(2);
            
            // 消除重力加速度作用
            Eigen::VectorXd g(3);
            g << 0, 0, gravity;
            Eigen::VectorXd da 
                = cpp_tools::Euler_Rotation_X(gamma).inverse()*
                cpp_tools::Euler_Rotation_Y(psi).inverse()*g;
            a = a - da;

            imu.setZero(5);
            imu(0) = psi;
            imu(1) = gamma;
            imu(2) = omega;
            imu(3) = a(0);
            imu(4) = a(1);
        }

        Eigen::VectorXd uwb_front = *m_mtData.uwb_front.read().access(); // 前面的uwb
        Eigen::VectorXd uwb_rear = *m_mtData.uwb_rear.read().access(); // 后面的uwb

        if ((uwb_front.array().abs() > 1e3).any() || (uwb_front.array().isNaN()).any())
        {
            // std::cout << "UWB a 疑似有错误读数: "
            // << uwb_front.transpose() 
            // << std::endl;

            return -1;
        }
        if ((uwb_rear.array() > 1e3).any() || (uwb_rear.array().isNaN()).any())
        {
            // std::cout << "UWB b 疑似有错误读数: "
            // << uwb_rear.transpose() 
            // << std::endl;

            return -1;
        }

        uwb.setZero(3);
        uwb(0) = (uwb_front(0)*dx_rear - uwb_rear(0)*dx_front)/(dx_rear - dx_front);
        uwb(1) = (uwb_front(1)*dx_rear - uwb_rear(1)*dx_front)/(dx_rear - dx_front);
        uwb(2) = atan2(uwb_front(1) - uwb_rear(1),uwb_front(0) - uwb_rear(0));

        *m_mtData.uwb_front_update.write().access() = 0;
        *m_mtData.uwb_rear_update.write().access() = 0;

        return 0;
    };

public:
    estimator_UWB_IMU()
        : estimator_base()
    {
        std::cout << "state estimator start ..." << std::endl;

        // data init
        m_mtData.IMU_filter.write().access()->resize(3,1);
        m_mtData.AHRS_filter.write().access()->resize(3,1);
        *m_mtData.uwb_front_update.write().access() = 0;
        *m_mtData.uwb_rear_update.write().access() = 0;
        m_mtData.state.write().access()->resize(10);

        # ifdef INPUT_REAL
        cpp_tools::XmlParameterManager m_xml("device", host_path + settings_path + "/device/settings.xml");
        bool writeXML = false;
        if (!m_xml.getParameter("IMU", "fileName", IMU_device))
        {
            std::cout << "IMU设备路径使用使用默认值: " << IMU_device << std::endl;
            m_xml.setParameter("IMU", "fileName", IMU_device);

            writeXML = true;
        }
        if (!m_xml.getParameter("IMU", "bauds", IMU_bauds))
        {
            std::cout << "IMU波特率使用使用默认值: " << IMU_bauds << std::endl;
            m_xml.setParameter("IMU", "bauds", IMU_bauds);

            writeXML = true;
        }
        if (!m_xml.getParameter("UWB", "fileName_A", UWB_A_device))
        {
            std::cout << "UWB_A设备路径使用使用默认值: " << UWB_A_device << std::endl;
            m_xml.setParameter("UWB", "fileName_A", UWB_A_device);
            
            writeXML = true;
        }
        if (!m_xml.getParameter("UWB", "fileName_B", UWB_B_device))
        {
            std::cout << "UWB_B设备路径使用使用默认值: " << UWB_B_device << std::endl;
            m_xml.setParameter("UWB", "fileName_B", UWB_B_device);
            
            writeXML = true;
        }
        if (!m_xml.getParameter("UWB", "bauds", UWB_bauds))
        {
            std::cout << "UWB波特率使用使用默认值: " << UWB_bauds << std::endl;
            m_xml.setParameter("UWB", "bauds", UWB_bauds);
            
            writeXML = true;
        }
        if (writeXML)
        {
            m_xml.saveToFile(host_path + settings_path + "/device/settings.xml");
        }
        
        // =============================imu 端口读取线程=============================
        std::thread t_imu_port([&]() -> void{
            WHEELTEC_N100::driver d_imu(IMU_device.data(), IMU_bauds);

            Eigen::MatrixXd IMU_filterData; // 加速度采用均值滤波
            Eigen::VectorXd AHRS_filterData; // 陀螺仪采用低通滤波
            Eigen::VectorXd IMU_last(3); // 记录上一个数据
            homo_ocp::clock m_clock;
            while (isRunning)
            {
                auto r = d_imu.reciveData();
                if (r == 1)
                {
                    auto imu_data = d_imu.get_IMUdata();

                    // std::cout << "recive IMU data" << std::endl;

                    // 滤波
                    double a_x = imu_data->accelerometer_x;
                    double a_y = imu_data->accelerometer_y;
                    double a_z = imu_data->accelerometer_z;
                    Eigen::VectorXd imuVector(3);
                    imuVector << a_x,a_y,a_z;
                    if ((imuVector.array().abs() > 16).any() || (imuVector.array().isNaN()).any() || imuVector.array().isZero())
                    {
                        // std::cout << "IMU 疑似有错误读数: " << imuVector.transpose() << std::endl;
                    }
                    else
                    {
                        if (IMU_filterData.size() == 0) // 初始化时全部给一样的
                        {
                            IMU_filterData.resize(avgFilterLen,3);

                            for (size_t i = 0; i < avgFilterLen; i++)
                            {
                                IMU_filterData.row(i) = imuVector.transpose();
                            }
                        }
                        else
                        {
                            // 异常值滤波：欧拉距离与前面的值均大于一定值则认为异常
                            int overDis = 0;
                            for (size_t i = 0; i < avgFilterLen; i++)
                            {
                                auto dis = IMU_filterData.row(i) - imuVector.transpose();
                                if ((dis.array().abs() > maxDis).any())
                                {
                                    overDis += 1;
                                }
                            }

                            // 检查上一个值与此次的值是否一致，如果一致则不认为是异常值
                            bool acceptThis = false;
                            if (((IMU_last - imuVector).array().abs() <= maxDis).all())
                            {
                                acceptThis = true;
                            }
                            
                            if (overDis < avgFilterLen || acceptThis)
                            {
                                Eigen::MatrixXd tmp = IMU_filterData.bottomRows(avgFilterLen - 1);
                                IMU_filterData.topRows(avgFilterLen - 1) = tmp;
                                IMU_filterData.bottomRows(1) = imuVector.transpose();
                            }

                            IMU_last = imuVector;
                        }

                        *m_mtData.IMU_filter.write().access() = IMU_filterData.colwise().sum()/avgFilterLen;
                    }
                }
                else if (r == 2)
                {
                    auto ahrs_data = d_imu.get_AHRSdata();

                    // std::cout << "recive AHRS data: " << ahrs_data->Heading << std::endl;

                    // 滤波
                    double psi = ahrs_data->Pitch;
                    double gamma = ahrs_data->Roll;
                    double omega = ahrs_data->HeadingSpeed;
                    Eigen::VectorXd ahrsVector(3);
                    ahrsVector << psi,gamma,omega;
                    if ((ahrsVector.array().abs() > 16).any() || (ahrsVector.array().isNaN()).any() || ahrsVector.array().isZero())
                    {
                        // std::cout << "AHRS 疑似有错误读数: " << ahrsVector.transpose() << std::endl;
                    }
                    else
                    {
                        if (AHRS_filterData.size() == 0) // 初始化时全部给一样的
                        {
                            AHRS_filterData = ahrsVector;
                        }
                        else
                        {
                            AHRS_filterData = AHRS_filterData + ahrsFilterFactor*(ahrsVector - AHRS_filterData);
                        }

                        *m_mtData.AHRS_filter.write().access() = AHRS_filterData;
                    }
                }

                std::this_thread::sleep_for(std::chrono::microseconds(std::max((int)(1e6/m_frequency.imu_port),1)));
            }
        });
        m_threadPool.push_back(t_imu_port);

        // =============================uwb端口读取线程=============================
        std::thread t_uwb_port([&]() -> void{
            LTP_A::tag d_uwb_A(UWB_A_device.data(), UWB_bauds);
            LTP_A::tag d_uwb_B(UWB_B_device.data(), UWB_bauds);

            auto t = std::chrono::steady_clock::now();
            while (isRunning)
            {
                if (d_uwb_A.reciveData() >= 1)
                {
                    auto uwb_A_data = d_uwb_A.get_data();

                    if (uwb_A_data->pos_3d)
                    {
                        // std::cout << "recive uwb tag A data" << std::endl;
                        // std::cout << uwb_A_data->local_time << std::endl;
                        // std::cout << uwb_A_data->system_time << std::endl;

                        if (uwb_A_data->id == front_UWB_ID)
                        {
                            *m_mtData.uwb_front.write().access() << 
                                uwb_A_data->pos_3d[0] - m_mtData.position_cor.read().access()->first,
                                uwb_A_data->pos_3d[1] - m_mtData.position_cor.read().access()->second;
                            *m_mtData.uwb_front_update.write().access() = 1;
                        }
                        else
                        {
                            *m_mtData.uwb_rear.write().access() << 
                                uwb_A_data->pos_3d[0] - m_mtData.position_cor.read().access()->first,
                                uwb_A_data->pos_3d[1] - m_mtData.position_cor.read().access()->second;
                            *m_mtData.uwb_rear_update.write().access() = 1;
                        }
                    }
                }
                if (d_uwb_B.reciveData() >= 1)
                {
                    auto uwb_B_data = d_uwb_B.get_data();

                    if (uwb_B_data->pos_3d)
                    {
                        // std::cout<<"recive uwb tag B data" << std::endl;
                        // std::cout << uwb_B_data->system_time << std::endl;

                        if (uwb_B_data->id == front_UWB_ID)
                        {
                            *m_mtData.uwb_front.write().access() << 
                                uwb_B_data->pos_3d[0] - m_mtData.position_cor.read().access()->first,
                                uwb_B_data->pos_3d[1] - m_mtData.position_cor.read().access()->second;
                            *m_mtData.uwb_front_update.write().access() = 1;
                        }
                        else
                        {
                            *m_mtData.uwb_rear.write().access() << 
                                uwb_B_data->pos_3d[0] - m_mtData.position_cor.read().access()->first,
                                uwb_B_data->pos_3d[1] - m_mtData.position_cor.read().access()->second;
                            *m_mtData.uwb_rear_update.write().access() = 1;
                        }
                    }
                }

                if (*m_mtData.uwb_front_update.write().access() == 1 && *m_mtData.uwb_rear_update.write().access() == 1)
                {
                    // std::cout<<"UWB 同步" << std::endl;
                    std::this_thread::sleep_until(t + std::chrono::microseconds(std::max((int)(1e6/m_frequency.uwb_port),1)));
                    t = std::chrono::steady_clock::now();
                }
                else
                {
                    // std::cout<<"UWB 未同步" << std::endl;
                    std::this_thread::sleep_for(std::chrono::microseconds(std::min((int)(1e6/m_frequency.uwb_port),10)));
                }
            }
        });
        m_threadPool.push_back(t_uwb_port);

        # else // INPUT_REAL
        // =============================从仿真环境中一次性获取数据=============================
        std::thread t_simState([&]() -> void {
            cpp_tools::udpVecReceiver<double> stateReceiver(udpSettings::PORT::lower_sim::state_len,udpSettings::PORT::lower_sim::state,1);

            cpp_tools::clock m_clock;
            while (isRunning)
            {
                auto t = std::chrono::steady_clock::now();

                int info = stateReceiver.receive();
                if (info == 0 && m_clock.toc(0) >= 10)
                {
                    m_clock.tic();

                    auto data = stateReceiver.get();
                    Eigen::VectorXd position_cor(2);
                    position_cor << m_mtData.position_cor.read().access()->first,m_mtData.position_cor.read().access()->second;

                    *m_mtData.uwb_front.write().access() = data.topRows(2) - position_cor;
                    *m_mtData.uwb_rear.write().access() = data.middleRows(2,2) - position_cor;
                    *m_mtData.IMU_filter.write().access() = data.middleRows(4,3);
                    *m_mtData.AHRS_filter.write().access() = data.bottomRows(3);

                    *m_mtData.uwb_front_update.write().access() = 1;
                    *m_mtData.uwb_rear_update.write().access() = 1;
                }
                
                std::this_thread::sleep_until(t + std::chrono::microseconds(std::max((int)(1e6/m_frequency.state),1)));
            }
        });
        m_threadPool.push_back(t_simState);

        # endif // INPUT_REAL
        
        // =============================状态估计线程=============================
        std::thread t_estimate_port([&]() -> void {
            // 问题初始化
            homo_ocp::prob m_prob;
            HOMO_OCP_PROB(m_prob,filter_W1M1);
            if (m_prob.n_x != nz_h[0] || m_prob.n_x != nz_A[1] || m_prob.n_eq != nz_A[0])
            {
                throw std::runtime_error("问题维度不一致，请检查提供问题及其梯度信息的文件");
            }

            homo_ocp::solver m_solver(m_prob,9999,1,-1);

            // 需要注意的是，这里的imu加速度是去除重力加速度后的结果
            Eigen::VectorXd imu = Eigen::VectorXd::Zero(5); // psi, gamma, omega_z, a_x, a_y
            Eigen::VectorXd uwb = Eigen::VectorXd::Zero(3); // x, y, varphi
            
            Eigen::VectorXd Q(3);
            Eigen::VectorXd R(5);
            Eigen::VectorXd D(1);
            
            // 循环计时器，用来记录两次uwb数据之间的间隔时间(用于差分计算速度)
            cpp_tools::clock loop_clock;
            loop_clock.tic();

            // 记录
            cpp_tools::csvRecorder timeData(host_path + data_path + "/estimate/time.csv"); // 包括当前时刻，循环时间，求解时间
            cpp_tools::csvRecorder imuData(host_path + data_path + "/estimate/imu.csv");
            cpp_tools::csvRecorder uwbData(host_path + data_path + "/estimate/uwb.csv");
            
            // 
            Eigen::VectorXd opt_X(5);
            auto t_loop = std::chrono::steady_clock::now();
            while (isRunning)
            {
                double opt_v_x = opt_X(3);
                double opt_v_y = opt_X(4);
                double opt_v = std::sqrt(std::pow(opt_v_x,2) + std::pow(opt_v_y,2));

                // 尝试读取数据
                if (getData(imu,uwb) < 0)
                {
                    std::this_thread::sleep_for(std::chrono::microseconds(100)); // 若失败则等待较短时间后重新读取

                    continue;
                }

                double t_read = cpp_tools::getCurrentTimeAsDouble();

                double psi = imu(0);
                double gamma = imu(1);
                double omega = imu(2);
                double a_x = imu(3);
                double a_y = imu(4);

                double x = uwb(0);
                double y = uwb(1);
                double varphi = uwb(2);

                // 因为引入了航向角的运动学模型，因此需要保证mod2pi不会导致数值突变
                while (std::abs(varphi - opt_X(2)) > pi)
                {
                    if (varphi - opt_X(2) > pi)
                    {
                        opt_X(2) += 2*pi;
                    }
                    else
                    {
                        opt_X(2) -= 2*pi;
                    }
                }
                
                // 当程序开始时，为了使速度估计能够快速收敛，会对QR权重进行调整
                double v_std = std::max(opt_v - 0.5,1.);
                Q << 4,4,1;
                R << 40,40,10,0.5/std::min(std::pow(v_std,3),10.),0.5/std::min(std::pow(v_std,3),10.);
                D << 10/std::min(std::pow(v_std,2),10.);

                double diff_omega = std::abs(omega) + 0.4;
                
                // 问题的初值
                Eigen::VectorXd r_0(m_prob.n_x);
                r_0 << opt_X, Eigen::VectorXd::Zero(5);

                // 计算数据的间隔时间，填充问题并求解
                double Delta_t = loop_clock.toc(0); // 单位 ms
                m_prob.p << 
                    Delta_t/1e3, // 参数中的时间单位是 s
                    opt_X,
                    uwb,
                    imu,
                    Q,
                    R,
                    D;
                loop_clock.tic();

                m_prob.lbx << x - 0.2, y - 0.2, -9999, -9999, -9999, -9999, -9999, -Delta_t/1e3*diff_omega, -9999, -9999;
                m_prob.ubx << x + 0.2, y + 0.2, 9999, 9999, 9999, 9999, 9999, Delta_t/1e3*diff_omega, 9999, 9999;

                m_solver.reLoadBoundary(&m_prob.lbx,&m_prob.ubx,0);
                m_solver.reLoadParam(m_prob.p);

                auto sol = m_solver.solve(r_0);
                auto cpuTime = m_solver.getCpuTime();

                opt_X = sol.x.topRows(5);
                while (std::abs(opt_X(2)) > pi)
                {
                    if (opt_X(2) > pi)
                    {
                        opt_X(2) -= 2*pi;
                    }
                    else
                    {
                        opt_X(2) += 2*pi;
                    }
                }

                Eigen::VectorXd out(10);
                out << opt_X(0), opt_X(1), opt_X(2), psi, gamma, opt_X(3), opt_X(4), omega, a_x, a_y;
                *(m_mtData.state.write().access()) = out;

                // 打印和输出
                // {
                //     int i = 0;
                //     std::cout << "Estimator: "
                //     << "\n  x: " << out(i++)
                //     << "\n  y: " << out(i++)
                //     << "\n  varphi: " << out(i++)
                //     << "\n  psi: " << out(i++)
                //     << "\n  gamma: " << out(i++)
                //     << "\n  v_x: " << out(i++)
                //     << "\n  v_y: " << out(i++)
                //     << "\n  omega: " << out(i++)
                //     << "\n  a_x: " << out(i++)
                //     << "\n  a_y: " << out(i++)
                //     << std::endl;
                // }

                t_update = t_read;

                // 数据记录
                Eigen::VectorXd time(3);
                time << cpp_tools::getCurrentTimeAsDouble(), Delta_t, cpuTime;
                timeData.record(time);
                imuData.record(imu);
                uwbData.record(uwb);

                // loop
                std::this_thread::sleep_until(t_loop + std::chrono::microseconds(std::max((int)(1e6/m_frequency.estimate),1)));
                t_loop = std::chrono::steady_clock::now();
            }
        });
        m_threadPool.push_back(t_estimate_port);
    };

    ~estimator_UWB_IMU()
    {
        shutdown();

        std::cout << "state estimator exit." << std::endl;
    };

    void shutdown()  override
    {
        isRunning = false;

        m_threadPool.joinAll();
    };

    Eigen::VectorXd getStateAsVector()  override
    {
        auto state = *m_mtData.state.read().access();

        return state;
    }

    state getStateAsStruct() override
    {
        auto stateVec = getStateAsVector();

        int i = 0;
        return {
            stateVec(i++),
            stateVec(i++),
            stateVec(i++),
            stateVec(i++),
            stateVec(i++),
            stateVec(i++),
            stateVec(i++),
            stateVec(i++),
            stateVec(i++),
            stateVec(i++)
        };
    };

    /// @brief 获取最近一次状态更新的时间，-1表示没有任何更新
    /// @return 
    double getUpdateTime() override
    {
        return t_update;
    }

    /// @brief 增加uwb定位x和y的修正量，偏大给正值
    /// @param x 
    /// @param y 
    void reInitFrame(double x, double y) override
    {
        m_mtData.position_cor.write().access()->first += x;
        m_mtData.position_cor.write().access()->second += y;
    }
};
